#include <os/driver.h>
#include <os/initcall.h>
#include <os/hardirq.h>
#include <os/virmem.h>
#include <os/diskman.h>
#include <arch/pci.h>
#include <arch/vmm.h>
#include <arch/x86.h>
#include <arch/memory.h>
#include <driver/ahci.h>
#include <lib/type.h>
#include <lib/stdio.h>
#include <lib/assert.h>
#include <lib/string.h>
#include <sys/ioctl.h>
#include <driver/ahci.h>

static pci_dev_t *ahci_pci;
static int ahci_int = 0;
static struct hba_memory *hba_mem;
static device_extension_t *ports[32];
static int ahci_next_device = 0;

pci_dev_t *get_ahci_pci(void)
{
	pci_dev_t *ahci = PciGetDeviceByClass(0x1, 0x6);
	if (!ahci)
		ahci = PciGetDeviceByClass(0x8086, 0x8c03);
	if (!ahci)
		ahci = PciGetDeviceByClass(0x8086, 0x2922);
	if (!ahci)
		return NULL;
#ifdef DEBUG_AHCI
	dbgprint("[ahci]: device vendorID %x deviceID %x class code %x\n", ahci->vendor_id, ahci->device_id, ahci->class_code);
#endif
	// pci_device_dump(ahci);

	PciEnableBusMaster(ahci);

	/* 映射IO物理内存地址到虚拟地址中，才能对设备映射到内存的地址进行操作 */
	hba_mem = MemIoReMap((address_t)ahci->bar[5].base, ahci->bar[5].len);
	if (hba_mem == NULL)
	{
		KPrint("[ahci] device memio_remap on %x length %x failed!\n", ahci->bar[5].base, ahci->bar[5].len);
		return NULL;
	}
	tlb_flush(); // 刷新快表
#ifdef DEBUG_AHCI
	dbgprint("[ahci]: mapping hba_mem to %x -> %x\n", hba_mem, ahci->bar[5].base_addr);
	dbgprint("[ahci]: using interrupt %d\n", ahci->irq_line);
#endif
	ahci_int = ahci->irq;

	return ahci;
}

uint32_t ahci_flush_commands(struct hba_port *port)
{
	/* the commands may not take effect until the command
	 * register is read again by software, because reasons.
	 */
	volatile uint32_t c = port->command;
	c = c;
	return c;
}

void ahci_stop_port_command_engine(volatile struct hba_port *port)
{
	port->command &= ~HBA_PxCMD_ST;
	port->command &= ~HBA_PxCMD_FRE;
	while ((port->command & HBA_PxCMD_CR) || (port->command & HBA_PxCMD_FR))
		CpuPause();
}

void ahci_start_port_command_engine(volatile struct hba_port *port)
{
	while (port->command & HBA_PxCMD_CR)
		;
	port->command |= HBA_PxCMD_FRE;
	port->command |= HBA_PxCMD_ST;
	ahci_flush_commands((struct hba_port *)port);
}

void ahci_init_hba(struct hba_memory *abar)
{
	if (abar->ext_capabilities & 1)
	{
		/* request BIOS/OS ownership handoff */
		KPrint(PRINT_NOTICE "[ahci]: requesting AHCI ownership change\n");
		abar->bohc |= (1 << 1);
		while ((abar->bohc & 1) || !(abar->bohc & (1 << 1)))
			CpuPause();
		KPrint(PRINT_NOTICE "[ahci]: ownership change completed\n");
	}

	/* enable the AHCI and reset it */
	abar->global_host_control |= HBA_GHC_AHCI_ENABLE;
	abar->global_host_control |= HBA_GHC_RESET;
	/* wait for reset to complete */
	while (abar->global_host_control & HBA_GHC_RESET)
		CpuPause();
	/* enable the AHCI and interrupts */
	abar->global_host_control |= HBA_GHC_AHCI_ENABLE;
	abar->global_host_control |= HBA_GHC_INTERRUPT_ENABLE;
	Mdelay(20);
#ifdef DEBUG_AHCI
	KPrint(PRINT_INFO "[ahci]: caps: %x %x ver:%x ctl: %x\n", abar->capability, abar->ext_capabilities, abar->version, abar->global_host_control);
#endif
}

struct hba_command_header *ahci_initialize_command_header(struct hba_memory *abar, struct hba_port *port, device_extension_t *dev, int slot, int write, int atapi, int prd_entries, int fis_len)
{
	struct hba_command_header *h = (struct hba_command_header *)dev->clb_virt;
	h += slot;
	h->write = write ? 1 : 0;
	h->prdb_count = 0;
	h->atapi = atapi ? 1 : 0;
	h->fis_length = fis_len;
	h->prdt_len = prd_entries;
	h->prefetchable = 0;
	h->bist = 0;
	h->pmport = 0;
	h->reset = 0;
	return h;
}

struct fis_reg_host_to_device *ahci_initialize_fis_host_to_device(struct hba_memory *abar, struct hba_port *port, device_extension_t *dev, int slot, int cmdctl, int ata_command)
{
	struct hba_command_table *tbl = (struct hba_command_table *)(dev->ch[slot]);
	struct fis_reg_host_to_device *fis = (struct fis_reg_host_to_device *)(tbl->command_fis);

	memset(fis, 0, sizeof(*fis));
	fis->fis_type = FIS_TYPE_REG_H2D;
	fis->command = ata_command;
	fis->c = cmdctl ? 1 : 0;
	return fis;
}

void ahci_send_command(struct hba_port *port, int slot)
{
	port->interrupt_status = ~0;
	port->command_issue = (1 << slot);
	ahci_flush_commands(port);
}

int ahci_write_prdt(struct hba_memory *abar, struct hba_port *port, device_extension_t *dev, int slot, int offset, int length, address_t virt_buffer)
{
	int num_entries = ((length - 1) / PRDT_MAX_COUNT) + 1;
	struct hba_command_table *tbl = (struct hba_command_table *)(dev->ch[slot]);
	int i;
	struct hba_prdt_entry *prd;
	for (i = 0; i < num_entries - 1; i++)
	{
		/* TODO: do we need to do this? */
		address_t phys_buffer;
		phys_buffer = Vbase2Pybase(virt_buffer);
		// mm_virtual_getmap(virt_buffer, &phys_buffer, 0);
		prd = &tbl->prdt_entries[i + offset];
		prd->byte_count = PRDT_MAX_COUNT - 1;
		prd->data_base_l = phys_buffer & 0xffffffff;
		prd->data_base_h = 0;
		prd->interrupt_on_complete = 0;

		length -= PRDT_MAX_COUNT;
		virt_buffer += PRDT_MAX_COUNT;
	}
	address_t phys_buffer;
	phys_buffer = Vbase2Pybase(virt_buffer);
	// mm_virtual_getmap(virt_buffer, &phys_buffer, 0);
	prd = &tbl->prdt_entries[i + offset];
	prd->byte_count = length - 1;
	prd->data_base_l = phys_buffer & 0xffffffff;
	prd->data_base_h = 0;
	prd->interrupt_on_complete = 0;

	return num_entries;
}

void ahci_reset_device(struct hba_memory *abar, struct hba_port *port, device_extension_t *dev)
{
/* TODO: This needs to clear out old commands and lock properly so that new commands can't get sent
 * while the device is resetting */
#ifdef DEBUG_AHCI
	KPrint(PRINT_NOTICE "[ahci]: device %d: sending COMRESET and reinitializing\n", dev->idx);
#endif
	ahci_stop_port_command_engine(port);
	port->sata_error = ~0;
	/* power on, spin up */
	port->command |= 2;
	port->command |= 4;
	ahci_flush_commands(port);
	Mdelay(1);
	/* initialize state */
	port->interrupt_status = ~0;			   /* clear pending interrupts */
	port->interrupt_enable = AHCI_DEFAULT_INT; /* we want some interrupts */
	port->command &= ~((1 << 27) | (1 << 26)); /* clear some bits */
	port->sata_control |= 1;
	Mdelay(10);
	port->sata_control |= (~1);
	Mdelay(10);
	port->interrupt_status = ~0;			   /* clear pending interrupts */
	port->interrupt_enable = AHCI_DEFAULT_INT; /* we want some interrupts */
	ahci_start_port_command_engine(port);
	dev->solts = 0;
	port->sata_error = ~0;
}

int ahci_port_dma_data_transfer(struct hba_memory *abar, struct hba_port *port, device_extension_t *dev, int slot, int write, address_t virt_buffer, int sectors, uint64_t lba)
{
	int timeout;
	int fis_len = sizeof(struct fis_reg_host_to_device) / 4;
	int ne = ahci_write_prdt(abar, port, dev,
							 slot, 0, ATA_SECTOR_SIZE * sectors, virt_buffer);
	ahci_initialize_command_header(abar,
								   port, dev, slot, write, 0, ne, fis_len);
	struct fis_reg_host_to_device *fis = ahci_initialize_fis_host_to_device(abar,
																			port, dev, slot, 1, write ? ATA_CMD_WRITE_DMA_EX : ATA_CMD_READ_DMA_EX);
	fis->device = 1 << 6;
	/* WARNING: assumes little-endian */
	fis->count_l = sectors & 0xFF;
	fis->count_h = (sectors >> 8) & 0xFF;

	fis->lba0 = (unsigned char)(lba & 0xFF);
	fis->lba1 = (unsigned char)((lba >> 8) & 0xFF);
	fis->lba2 = (unsigned char)((lba >> 16) & 0xFF);
	fis->lba3 = (unsigned char)((lba >> 24) & 0xFF);
	fis->lba4 = (unsigned char)((lba >> 32) & 0xFF);
	fis->lba5 = (unsigned char)((lba >> 40) & 0xFF);
	port->sata_error = ~0;
	timeout = ATA_TFD_TIMEOUT;
	while ((port->task_file_data & (ATA_DEV_BUSY | ATA_DEV_DRQ)) && --timeout)
	{
		// tm_schedule();
		//  cpu yield
		CpuPause();
		// task_yield();
	}
	if (!timeout)
		goto port_hung;

	port->sata_error = ~0;
	ahci_send_command(port, slot);
	timeout = ATA_TFD_TIMEOUT;
	while ((port->task_file_data & (ATA_DEV_BUSY | ATA_DEV_DRQ)) && --timeout)
	{
		CpuPause();
		// task_yield();
	}
	if (!timeout)
		goto port_hung;

	timeout = AHCI_CMD_TIMEOUT;
	while (--timeout)
	{
		if (!((port->sata_active | port->command_issue) & (1 << slot)))
			break;
		CpuPause();
		// task_yield();
	}
	if (!timeout)
		goto port_hung;
	if (port->sata_error)
	{
		KPrint(PRINT_ERR "[ahci]: device %d: ahci error\n", dev->idx);
		goto error;
	}
	if (port->task_file_data & ATA_DEV_ERR)
	{
		KPrint(PRINT_ERR "[ahci]: device %d: task file data error\n", dev->idx);
		goto error;
	}
	return 1;
port_hung:
	KPrint(PRINT_ERR "[ahci]: device %d: port hung\n", dev->idx);
error:
	KPrint(PRINT_ERR "[ahci]: device %d: tfd=%x, serr=%x\n",
		   dev->idx, port->task_file_data, port->sata_error);
	ahci_reset_device(abar, port, dev);
	return 0;
}

int ahci_device_identify_ahci(struct hba_memory *abar,
							  struct hba_port *port, device_extension_t *dev)
{
	int fis_len = sizeof(struct fis_reg_host_to_device) / 4;
	struct dma_region dma;
	dma.p.size = 0x1000;
	dma.p.alignment = 0x1000;
	dma.flags = DMA_REGION_SPECIAL;
	DmaAllocBuffer(&dma);
	ahci_write_prdt(abar, port, dev, 0, 0, 512, (address_t)dma.v);
	ahci_initialize_command_header(abar,
								   port, dev, 0, 0, 0, 1, fis_len);
	ahci_initialize_fis_host_to_device(abar,
									   port, dev, 0, 1, ATA_CMD_IDENITFY);
	int timeout = ATA_TFD_TIMEOUT;
	port->sata_error = ~0;
	while ((port->task_file_data & (ATA_DEV_BUSY | ATA_DEV_DRQ)) && --timeout)
		CpuPause();
	if (!timeout)
	{
		KPrint(PRINT_ERR "[ahci]: device %d: identify 1: port hung\n", dev->idx);
		KPrint(PRINT_ERR "[ahci]: device %d: identify 1: tfd=%x, serr=%x\n",
			   dev->idx, port->task_file_data, port->sata_error);
		DmaFreeBuffer(&dma);
		return 0;
	}

	ahci_send_command(port, 0);
#ifdef DEBUG_AHCI
	KPrint(PRINT_DEBUG "[AHCI]: port %d sata active %x command issue %x\n", dev->idx, port->sata_active, port->command_issue);
#endif
	timeout = AHCI_CMD_TIMEOUT;
	while (--timeout)
	{
		if (!((port->sata_active | port->command_issue) & 1))
			break;
	}
#ifdef DEBUG_AHCI
	KPrint(PRINT_DEBUG "[AHCI]: port %d sata active %x command issue %x\n", dev->idx, port->sata_active, port->command_issue);
#endif
	if (!timeout)
	{
		KPrint(PRINT_ERR "[ahci]: device %d: identify 2: port hung\n", dev->idx);
		KPrint(PRINT_ERR "[ahci]: device %d: identify 2: tfd=%x, serr=%x\n",
			   dev->idx, port->task_file_data, port->sata_error);
		DmaFreeBuffer(&dma);
		return 0;
	}

	memcpy(&dev->identify, (void *)dma.v, sizeof(struct ata_identify));
	DmaFreeBuffer(&dma);
#ifdef DEBUG_AHCI
	KPrint(PRINT_INFO "[ahci]: device %d: num sectors=%d: %x\n", dev->idx,
		   dev->identify.lba48_addressable_sectors, dev->identify.ss_2);
#endif

	if (!dev->identify.lba48_addressable_sectors)
		return 0;
	return 1;
}

uint32_t ahci_check_type(volatile struct hba_port *port)
{
	port->command &= ~1;
	while (port->command & (1 << 15))
		CpuPause();
	port->command &= ~(1 << 4);
	while (port->command & (1 << 14))
		CpuPause();
	wmb();
	port->command |= 2;
	wmb();
	Mdelay(10);

	uint32_t s = port->sata_status;

	// KPrint(PRINT_INFO "[ahci]: port data: sig=%x, stat=%x, ctl=%x, sac=%x\n", port->signature, s, port->command, port->sata_active);
	uint8_t ipm, det;
	ipm = (s >> 8) & 0x0F;
	det = s & 0x0F;
	// KPrint(PRINT_INFO "[ahci]: port check: ipm=%x, det=%x\n", ipm, det);
	if (ipm != HBA_PORT_IPM_ACTIVE || det != HBA_PORT_DET_PRESENT)
		return AHCI_DEV_NULL;
	switch (port->signature)
	{
	case SATA_SIG_ATAPI:
		return AHCI_DEV_SATAPI;
	case SATA_SIG_SEMB:
		return AHCI_DEV_SEMB;
	case SATA_SIG_PM:
		return AHCI_DEV_PM;
	default:
		return AHCI_DEV_SATA;
	}
	return AHCI_DEV_SATA;
}

int ahci_initialize_device(struct hba_memory *abar, device_extension_t *dev)
{
#ifdef DEBUG_AHCI
	KPrint(PRINT_INFO "[ahci]: initializing device %d\n", dev->idx);
#endif
	struct hba_port *port = (struct hba_port *)&abar->ports[dev->idx];
	ahci_stop_port_command_engine(port);
	port->sata_error = ~0;
	/* power on, spin up */
	port->command |= (2 | 4);
	ahci_flush_commands(port);
	Mdelay(2);
	/* initialize state */
	port->interrupt_status = ~0;			   /* clear pending interrupts */
	port->interrupt_enable = AHCI_DEFAULT_INT; /* we want some interrupts */

	port->command &= ~1;
	while (port->command & (1 << 15))
		CpuPause();
	port->command &= ~((1 << 27) | (1 << 26) | 1); /* clear some bits */
	ahci_flush_commands(port);
#ifdef DEBUG_AHCI
	dbgprint("[AHCI]: step1: port %d sata status %x control %x.\n", dev->idx,
			 port->sata_status, port->sata_control);
#endif
	/* start reset sata */
	port->sata_control |= 1;
	Mdelay(20);
#ifdef DEBUG_AHCI
	dbgprint("[AHCI]: step2: port %d sata status %x control %x.\n", dev->idx,
			 port->sata_status, port->sata_control);
#endif
	/* close DET, after init sata device done. */
	port->sata_control &= (~1);
	Mdelay(10);
#ifdef DEBUG_AHCI
	dbgprint("[AHCI]: step3: port %d sata status %x control %x.\n", dev->idx,
			 port->sata_status, port->sata_control);
#endif
	while (!(port->sata_status & 1))
		CpuPause();
	port->sata_error = ~0;
	port->command |= (1 << 28); /* set interface to active */
	while ((port->sata_status >> 8) != 1)
		CpuPause();
	port->interrupt_status = ~0;			   /* clear pending interrupts */
	port->interrupt_enable = AHCI_DEFAULT_INT; /* we want some interrupts */
#ifdef DEBUG_AHCI
	dbgprint("[AHCI]: map command list dma addr and fis dma addr start.\n");
#endif
	/* map memory */
	address_t clb_phys, fis_phys;

	dev->dma_clb.p.size = 0x2000;
	dev->dma_clb.p.alignment = 0x1000;
	dev->dma_fis.p.size = 0x1000;
	dev->dma_fis.p.alignment = 0x1000;

	dev->dma_clb.flags = DMA_REGION_SPECIAL;
	dev->dma_fis.flags = DMA_REGION_SPECIAL;
	DmaAllocBuffer(&dev->dma_clb);
	DmaAllocBuffer(&dev->dma_fis);

	dev->clb_virt = (void *)dev->dma_clb.v;
	dev->fis_virt = (void *)dev->dma_fis.v;
	clb_phys = dev->dma_clb.p.addr;
	fis_phys = dev->dma_fis.p.addr;
	dev->solts = 0;
	struct hba_command_header *h = (struct hba_command_header *)dev->clb_virt;
	int i;
	for (i = 0; i < HBA_COMMAND_HEADER_NUM; i++)
	{
		dev->ch_dmas[i].p.size = 0x1000;
		dev->ch_dmas[i].p.alignment = 0x1000;
		dev->ch_dmas[i].flags = DMA_REGION_SPECIAL;
		DmaAllocBuffer(&dev->ch_dmas[i]);
		dev->ch[i] = (void *)dev->ch_dmas[i].v;
		memset(h, 0, sizeof(*h));
		h->command_table_base_l = (dev->ch_dmas[i].p.addr & 0xffffffff);
		h->command_table_base_h = 0;
		h++;
	}

	port->command_list_base_l = (clb_phys & 0xffffffff);
	port->command_list_base_h = 0;

	port->fis_base_l = (fis_phys & 0xffffffff);
	port->fis_base_h = 0;
	ahci_start_port_command_engine(port);
	port->sata_error = ~0;
#ifdef DEBUG_AHCI
	dbgprint("[AHCI]: map command list dma addr and fis dma addr done.\n");
#endif
	return ahci_device_identify_ahci(abar, port, dev);
}

iostatus_t ahci_create_device(driver_object_t *driver, device_extension_t *dev)
{
	iostatus_t status;
	device_object_t *devobj;
	char devname[DEVICE_NAME_LEN+1] = {0};

	sprintf(devname, "%s%c", DEVICE_NAME, 'a' + ahci_next_device);
	ahci_next_device++;
	/* 初始化一些其它内容 */
	status = IoCreateDevice(driver, 0, devname, DEVICE_TYPE_DISK, &devobj);
	if (status != IO_SUCCESS)
	{
		KPrint(PRINT_ERR "[ahci]: create device on port %d failed!\n", dev->idx);
		return status;
	}
	/* buffered io mode */
	devobj->flags = DEVICE_BUFFER_IO;
	devobj->device_extension = dev;
	dev->device_object = devobj;
	dev->created = 1;
	dev->rwoffset = 0;
	return status;
}

int ahci_probe_ports(driver_object_t *driver, struct hba_memory *abar)
{
	uint32_t pi = abar->port_implemented;
#ifdef DEBUG_AHCI
	KPrint(PRINT_DEBUG "[ahci]: ports implemented: %x\n", pi);
#endif
	int counts = 0; /* exist device count */
	int i = 0;
	while (i < 32)
	{
		if (pi & 1)
		{
			uint32_t type = ahci_check_type(&abar->ports[i]);
			if (type == AHCI_DEV_SATA)
			{ /* SATA device */
#ifndef CONFIG_LIVECD
#ifdef DEBUG_AHCI
				KPrint(PRINT_DEBUG "[ahci]: detected SATA device on port %d\n", i);
#endif
				/* 创建设备扩展 */
				ports[i] = KMemAlloc(sizeof(device_extension_t));
				ports[i]->type = type;
				ports[i]->idx = i;
				MutexlockInit(&(ports[i]->lock));
				if (ahci_initialize_device(abar, ports[i]))
				{
/* create one device on port i */
#ifdef DEBUG_AHCI
					KPrint(PRINT_DEBUG "[ahci]: success to initialize device %d, disabling port\n", i);
#endif
					if (ahci_create_device(driver, ports[i]) < 0)
					{
						KPrint(PRINT_ERR "[ahci]: failed to create device %d, disabling port\n", i);
					}
					counts++;
				}
				else
				{
					KPrint(PRINT_ERR "[ahci]: failed to initialize device %d, disabling port\n", i);
				}
#else
				KPrint(PRINT_WARNNING "[ahci]: not support SATA device on port %d now!\n", i);
#endif
			}
			else if (type == AHCI_DEV_SATAPI)
			{ /* SATA device */
				KPrint(PRINT_WARNNING "[ahci]: not support SATAPI device on port %d now!\n", i);
			}
			else if (type == AHCI_DEV_PM)
			{ /* SATA device */
				KPrint(PRINT_WARNNING "[ahci]: not support Port multiplier on port %d now!\n", i);
			}
			else if (type == AHCI_DEV_SEMB)
			{ /* SATA device */
				KPrint(PRINT_WARNNING "[ahci]: not support Enclosure management bridge on port %d now!\n", i);
			}
			/* 暂时不处理其它类型的设备 */
		}
		i++;
		pi >>= 1;
	}
	return counts;
}

int ahci_port_acquire_slot(device_extension_t *dev)
{
	while (1)
	{
		int i;
		MutexlockLock(&dev->lock, MUTEX_LOCK_MODE_BLOCK);
		for (i = 0; i < 32; i++)
		{
			if (!(dev->solts & (1 << i)))
			{
				dev->solts |= (1 << i);
				MutexlockUnlock(&dev->lock);
				return i;
			}
		}
		MutexlockUnlock(&dev->lock);
		// yield
		CpuPause();
		// task_yield();
	}
}

void ahci_port_release_slot(device_extension_t *dev, int slot)
{
	MutexlockLock(&dev->lock, MUTEX_LOCK_MODE_BLOCK);
	dev->solts &= ~(1 << slot);
	MutexlockUnlock(&dev->lock);
}

/* since a DMA transfer must write to contiguous physical RAM, we need to allocate
 * buffers that allow us to create PRDT entries that do not cross a page boundary.
 * That means that each PRDT entry can transfer a maximum of PAGE_SIZE bytes (for
 * 0x1000 page size, that's 8 sectors). Thus, we allocate a buffer that is page aligned,
 * in a multiple of PAGE_SIZE, so that the PRDT will write to contiguous physical ram
 * (the key here is that the buffer need not be contiguous across multiple PRDT entries).
 */
int ahci_rw_multiple_do(int rw, int min, uint64_t blk, unsigned char *out_buffer, int count)
{
	uint32_t length = count * ATA_SECTOR_SIZE;
	int d = min;
	device_extension_t *dev = ports[d];
	uint64_t end_blk = dev->identify.lba48_addressable_sectors;
	if (blk >= end_blk)
	{
		KPrint("ahci: lba %d out of range %d\n", blk, end_blk);
		return 0;
	}

	if ((blk + count) > end_blk)
		count = end_blk - blk;
	if (!count)
		return 0;
	int num_pages = ((ATA_SECTOR_SIZE * (count - 1)) / PAGE_SIZE) + 1;
	assert(length <= (unsigned)num_pages * 0x1000);
	struct dma_region dma;
	dma.p.size = 0x1000 * num_pages;
	dma.p.alignment = 0x1000;
	dma.flags = DMA_REGION_SPECIAL;
	DmaAllocBuffer(&dma);
	int num_read_blocks = count;
	struct hba_port *port = (struct hba_port *)&hba_mem->ports[dev->idx];
	if (rw == 1)
		memcpy((void *)dma.v, out_buffer, length);

	int slot = ahci_port_acquire_slot(dev);
	if (!ahci_port_dma_data_transfer(hba_mem, port, dev, slot, rw == 1 ? 1 : 0, (address_t)dma.v, count, blk))
		num_read_blocks = 0;

	ahci_port_release_slot(dev, slot);

	if (rw == 0 && num_read_blocks)
	{
		memcpy(out_buffer, (void *)dma.v, length);
	}

	DmaFreeBuffer(&dma);
	return num_read_blocks * ATA_SECTOR_SIZE;
}

/* and then since there is a maximum transfer amount because of the page size
 * limit, wrap the transfer function to allow for bigger transfers than that even.
 */
int ahci_rw_multiple(int rw, int min, uint64_t blk, unsigned char *out_buffer, int count)
{
	int i = 0;
	int ret = 0;
	int c = count;
	for (i = 0; i < count; i += (PRDT_MAX_ENTRIES * PRDT_MAX_COUNT) / ATA_SECTOR_SIZE)
	{
		int n = (PRDT_MAX_ENTRIES * PRDT_MAX_COUNT) / ATA_SECTOR_SIZE;
		if (n > c)
			n = c;
		ret += ahci_rw_multiple_do(rw, min, blk + i, out_buffer + ret, n);
		c -= n;
	}
	return ret;
}

static int ahci_read_sector(device_extension_t *ext,
							unsigned int lba,
							void *buf,
							unsigned int count)
{
	return ahci_rw_multiple(0, ext->idx, lba, buf, count);
}

static int ahci_write_sector(
	device_extension_t *ext,
	unsigned int lba,
	void *buf,
	unsigned int count)
{
	return ahci_rw_multiple(1, ext->idx, lba, buf, count);
}

iostatus_t ahci_devctl(device_object_t *device, io_request_t *ioreq)
{
	unsigned int ctlcode = ioreq->parame.devctl.code;
	unsigned long arg = ioreq->parame.devctl.arg;
	unsigned long off;
	device_extension_t *ext = device->device_extension;

	iostatus_t status = IO_SUCCESS;
	int infomation = 0;
	switch (ctlcode)
	{
	case DISKIO_GETSIZE:
		*((unsigned int *)arg) = ext->size;
		break;
	case DISKIO_CLEAN:
		break;
	case DISKIO_SETOFF:
		off = *((unsigned long *)arg);
		if (off > ext->size - 1)
			off = ext->size - 1;
		ext->rwoffset = off;
		break;
	case DISKIO_GETOFF:
		*((unsigned long *)arg) = ext->rwoffset;
	default:
		infomation = -1;
		status = IO_FAILED;
		break;
	}
	ioreq->io_status.status = status;
	ioreq->io_status.info = infomation;
	IoCompleteRequest(ioreq);
	return status;
}

iostatus_t ahci_read(device_object_t *device, io_request_t *ioreq)
{
	long len;
	iostatus_t status = IO_SUCCESS;
	sector_t sectors = DIV_ROUND_UP(ioreq->parame.read.len, SECTOR_SIZE);
	device_extension_t *ext = device->device_extension;

#ifdef DEBUG_AHCI
	KPrint(PRINT_DEBUG "ahci_read: buf=%x sectors=%d off=%x\n",
		   ioreq->system_buffer, sectors, ioreq->parame.read.offset);
#endif
	unsigned long off;
	if (ioreq->parame.read.offset == DISKOFF_MAX)
	{
		off = ext->rwoffset;
	}
	else
	{
		off = ioreq->parame.read.offset;
	}
	len = ahci_read_sector(device->device_extension, off,
						   ioreq->sys_buff, sectors);
	if (!len)
	{ /* 执行失败 */
		status = IO_FAILED;
		len = 0;
	}

	ioreq->io_status.status = status;
	ioreq->io_status.info = len;

	IoCompleteRequest(ioreq);

	return status;
}

iostatus_t ahci_write(device_object_t *device, io_request_t *ioreq)
{
	long len;
	iostatus_t status = IO_SUCCESS;
	sector_t sectors = DIV_ROUND_UP(ioreq->parame.write.len, SECTOR_SIZE);
	device_extension_t *ext = device->device_extension;

#ifdef DEBUG_AHCI
	KPrint(PRINT_DEBUG "ahci_write: buf=%x sectors=%d off=%x\n",
		   ioreq->system_buffer, sectors, ioreq->parame.write.offset);
#endif
	unsigned long off;
	if (ioreq->parame.write.offset == DISKOFF_MAX)
	{
		off = ext->rwoffset;
	}
	else
	{
		off = ioreq->parame.write.offset;
	}
	len = ahci_write_sector(device->device_extension, off,
							ioreq->sys_buff, sectors);

	if (!len)
	{ /* 执行失败 */
		status = IO_FAILED;
		len = 0;
	}

	ioreq->io_status.status = status;
	ioreq->io_status.info = len;

	IoCompleteRequest(ioreq);

	return status;
}

/**
 * ahci_handler - ahci硬盘中断处理函数
 * @irq: 中断号
 * @data: 中断的数据
 */
static int ahci_handler(irqno_t irq, void *data)
{
	int intrhandled = IRQ_NEXTONE;
	int i;
	for (i = 0; i < 32; i++)
	{
		if (hba_mem->interrupt_status & (1 << i))
		{
			KPrint("ahci: interrupt %d occur!\n", i);
			hba_mem->ports[i].interrupt_status = ~0;
			hba_mem->interrupt_status = (1 << i);
			ahci_flush_commands((struct hba_port *)&hba_mem->ports[i]);
			intrhandled = IRQ_HANDLE;
		}
	}
	return intrhandled;
}

static iostatus_t ahci_enter(driver_object_t *driver)
{
	iostatus_t status = IO_SUCCESS;
	KPrint(PRINT_INFO "[ahci]: initializing ahci driver...\n");
	if (!(ahci_pci = get_ahci_pci()))
	{
		KPrint(PRINT_ERR "[ahci]: no AHCI controllers present!\n");
		status = IO_FAILED;
		return status;
	}

	KPrint(PRINT_INFO "[ahci]: irqno %d\n", ahci_int);
	if (IrqRegister(ahci_int, ahci_handler, IRQ_SHARE, "ahci", "ahci driver", (void *)driver) < 0)
	{
		KPrint(PRINT_ERR "[ahci]: register interrupt failed!\n");
		/* 需要取消内存映射以及关闭ahci总线 */
		status = IO_FAILED;
		return status;
	}
	ahci_init_hba(hba_mem);
	if (!ahci_probe_ports(driver, hba_mem))
	{
		KPrint(PRINT_INFO "[ahci]: initializing ahci driver failed!.\n");
		IrqUnregister(ahci_int, (address_t *)driver);
		status = IO_FAILED;
		return status;
	}

	KPrint(PRINT_INFO "[ahci]: initializing ahci driver done.\n");
	return status;
}

static iostatus_t ahci_exit(driver_object_t *driver)
{
	/* 遍历所有对象 */
	device_object_t *devobj, *next;
	device_extension_t *ext;
	IrqUnregister(ahci_int, driver);

	/* 由于涉及到要释放devobj，所以需要使用safe版本 */
	list_traversal_all_owner_to_next_safe(devobj, next, &driver->device_list, list)
	{
		ext = devobj->device_extension;

		DmaFreeBuffer(&(ext->dma_clb));
		DmaFreeBuffer(&(ext->dma_fis));
		int j;
		for (j = 0; j < HBA_COMMAND_HEADER_NUM; j++)
			DmaFreeBuffer(&(ext->ch_dmas[j]));

		KMemFree(ext);
		IoDeleteDevice(devobj); /* 删除每一个设备 */
	}

	string_del(&driver->name); /* 删除驱动名 */
	return IO_SUCCESS;
}

iostatus_t ahci_driver_func(driver_object_t *driver)
{
	iostatus_t status = IO_SUCCESS;

	/* 绑定驱动信息 */
	driver->driver_enter = ahci_enter;
	driver->driver_exit = ahci_exit;

	driver->dispatch_fun[IOREQ_READ] = ahci_read;
	driver->dispatch_fun[IOREQ_WRITE] = ahci_write;
	driver->dispatch_fun[IOREQ_DEVCTL] = ahci_devctl;

	/* 初始化驱动名字 */
	string_new(&driver->name, DRIVER_NAME, DRIVER_NAME_LEN);
#ifdef DEBUG_AHCI
	KPrint(PRINT_DEBUG "ahci_driver_func: driver name=%s\n",
		   driver->name.text);
#endif
	return status;
}

static __init void ahci_driver_entry(void)
{
	/*if (DriverObjectCreate(ahci_driver_func) < 0)
	{
		KPrint(PRINT_ERR "[driver]: %s create driver failed!\n", __func__);
	}*/
}
driver_initcall(ahci_driver_entry);
